package com.demo.udf;

import com.demo.utils.IPUtils;
import org.apache.spark.broadcast.Broadcast;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.api.java.UDF1;

import java.io.Serializable;

public class SparkIpUDF implements Serializable {

    private final SparkSession spark;

    public SparkIpUDF(SparkSession spark) {
        this.spark = spark;
    }

    public void registerCommon() {
        // ipv4_to_num
        spark.udf().register("ipv4_to_num", new UDF1<String, Long>() {
            @Override
            public Long call(String ipv4) throws Exception {
                return IPUtils.ipv4ToNum(ipv4);
            }
        }, org.apache.spark.sql.types.DataTypes.LongType);

        // num_to_ipv4
        spark.udf().register("num_to_ipv4", new UDF1<Long, String>() {
            @Override
            public String call(Long num) throws Exception {
                return IPUtils.num2Ipv4(num);
            }
        }, org.apache.spark.sql.types.DataTypes.StringType);

        // ipv4_short
        spark.udf().register("ipv4_short", new UDF1<String, String>() {
            @Override
            public String call(String ipv4) throws Exception {
                return IPUtils.convertIpv4ToShort(ipv4);
            }
        }, org.apache.spark.sql.types.DataTypes.StringType);

        // ipv4_full
        spark.udf().register("ipv4_full", new UDF1<String, String>() {
            @Override
            public String call(String ipv4) throws Exception {
                return IPUtils.convertIpv4ToFull(ipv4);
            }
        }, org.apache.spark.sql.types.DataTypes.StringType);
    }

    /**
     * 1. 广播变量的值一定要在UDF函数的定义中获取
     * 2. 不能对广播变量的值做转换，比如字符串的截取
     * @param ipLibBD
     */
    public void registerIpSearch(final Broadcast<String[][]> ipLibBD) {
        spark.udf().register("ip_search", new UDF1<String, String>() {
            @Override
            public String call(String ip) throws Exception {
                String[][] ipLibArr = ipLibBD.value();
                int index = IPUtils.ipBinarySearch(ipLibArr, IPUtils.convertIpv4ToFull(ip));
                if (index >= 0) {
                    return ipLibArr[index][2];
                } else {
                    return "unknown";
                }
            }
        }, org.apache.spark.sql.types.DataTypes.StringType);
    }
}
